{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d70339c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import swifter\n",
    "import json\n",
    "from functools import partial\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.stats import wasserstein_distance\n",
    "import helpers as ph\n",
    "from termcolor import colored\n",
    "from helpers import PEW_SURVEY_LIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a249d20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_DIR = './data/human_resp/'\n",
    "RESULT_DIR = './data/runs'\n",
    "CONTEXT = \"default\" # One of [\"default\", \"steer-qa\", \"steer-bio\", \"steer-portray\"]\n",
    "OUTPUT_DIR = f'./data/distributions'\n",
    "if not os.path.exists(OUTPUT_DIR):\n",
    "    os.makedirs(OUTPUT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80175239",
   "metadata": {},
   "outputs": [],
   "source": [
    "if CONTEXT == \"default\":\n",
    "    SURVEY_LIST = [f'American_Trends_Panel_W{SURVEY_WAVE}' for SURVEY_WAVE in PEW_SURVEY_LIST] + \\\n",
    "                ['Pew_American_Trends_Panel_disagreement_500']\n",
    "else:\n",
    "    SURVEY_LIST = ['Pew_American_Trends_Panel_disagreement_500']\n",
    "    steer_df = pd.read_csv(f'./data/model_input/{CONTEXT}.csv',\n",
    "                       delimiter='\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c10ed27f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for SURVEY_NAME in SURVEY_LIST:\n",
    "    print(colored(SURVEY_NAME, \"red\"))\n",
    "    RESULT_FILES = [f for f in os.listdir(RESULT_DIR) if SURVEY_NAME in f and f'context={CONTEXT}' in f]\n",
    "    \n",
    "    ## Read human responses and survey info\n",
    "    info_df = pd.read_csv(os.path.join(DATASET_DIR, SURVEY_NAME, 'info.csv'))\n",
    "    info_df['option_ordinal'] = info_df.apply(lambda x: eval(x['option_ordinal']), axis=1)\n",
    "    info_df['references'] = info_df.apply(lambda x: eval(x['references']), axis=1)\n",
    "\n",
    "    md_df = pd.read_csv(os.path.join(DATASET_DIR, SURVEY_NAME, 'metadata.csv'))\n",
    "    md_df['options'] = md_df.apply(lambda x: eval(x['options']), axis=1)\n",
    "    md_order = {'Overall': {'Overall': 0}}\n",
    "    md_order.update({k: {o: oi for oi, o in enumerate(opts)} for k, opts in zip(md_df['key'], md_df['options'])})\n",
    "\n",
    "    ## Get model opinion distribution\n",
    "    print(colored('--Getting LM opinion distribution--', 'blue'))\n",
    "    model_df = ph.get_model_opinions(RESULT_DIR, RESULT_FILES, info_df)\n",
    "    \n",
    "    ## Get human opinion distribution\n",
    "    print(colored('--Getting human opinion distribution--', 'blue'))\n",
    "    if SURVEY_NAME != \"Pew_American_Trends_Panel_disagreement_500\":\n",
    "        resp_df = pd.read_csv(os.path.join(DATASET_DIR, SURVEY_NAME, 'responses.csv'))\n",
    "        human_df = pd.concat([ph.extract_human_opinions(resp_df, \n",
    "                                                        model_df, \n",
    "                                                        md_df, \n",
    "                                                        demographic=demographic, \n",
    "                                                        wave=int(SURVEY_NAME.split('_W')[1]))\n",
    "                   for demographic in ph.DEMOGRAPHIC_ATTRIBUTES])\n",
    "    else:\n",
    "        if CONTEXT != \"default\":\n",
    "            contexts = np.unique(model_df['context'])\n",
    "            steer_dict = ph.get_steering_group(CONTEXT, steer_df, contexts)\n",
    "            model_df['steer_attribute'] = model_df.apply(lambda x: steer_dict[x['context']]['attribute'], \n",
    "                                                 axis=1)\n",
    "            model_df['steer_group'] = model_df.apply(lambda x: steer_dict[x['context']]['group'], \n",
    "                                                         axis=1)\n",
    "            steer_groups = list(set(model_df['steer_group'].values)) + ['Overall']\n",
    "\n",
    "        human_df = []\n",
    "        for wave in PEW_SURVEY_LIST:\n",
    "            sn = f'American_Trends_Panel_W{wave}'\n",
    "            hdf = pd.read_csv(os.path.join(OUTPUT_DIR, f'{sn}_default_human.csv'))\n",
    "            idf = info_df[info_df['survey'] == f'Pew_{sn}']\n",
    "            hdf = hdf[hdf['qkey'].isin(idf['key'].values)]\n",
    "            human_df.append(hdf)\n",
    "        human_df = pd.concat(human_df)\n",
    "        if CONTEXT != \"default\": human_df = human_df[human_df['group'].isin(steer_groups)]\n",
    "        human_df['D_H'] = human_df.apply(lambda x: [float(f) for f in x['D_H'][1:-1].strip().split(' ') if len(f)], axis=1)\n",
    "        \n",
    "        \n",
    "    human_df.to_csv(os.path.join(OUTPUT_DIR, f'{SURVEY_NAME}_{CONTEXT}_human.csv'))\n",
    "    model_df.to_csv(os.path.join(OUTPUT_DIR, f'{SURVEY_NAME}_{CONTEXT}_model.csv'))\n",
    "\n",
    "    ## Combine and save\n",
    "    print(colored('--Comparing opinion distribution--', 'blue'))\n",
    "    combined_df = pd.merge(model_df, human_df)\n",
    "    combined_df['group_order'] = combined_df.apply(lambda x: md_order[x['attribute']][x['group']], axis=1)\n",
    "    combined_df['WD'] = combined_df.swifter.apply(lambda x: wasserstein_distance(x['ordinal'], \n",
    "                                                                         x['ordinal'],\n",
    "                                                                         x['D_M'], x['D_H']) / ph.get_max_wd(x['ordinal']), \n",
    "                                          axis=1)\n",
    "\n",
    "    combined_df.to_csv(os.path.join(OUTPUT_DIR, f'{SURVEY_NAME}_{CONTEXT}_combined.csv'))\n",
    "    \n",
    "    if CONTEXT == 'default':\n",
    "        print(colored('--Getting human baseline--', 'blue'))\n",
    "        human_overall = human_df[human_df['group'] == 'Overall'].rename(columns={'D_H': 'D_O', 'R_H': 'R_O'})\n",
    "        human_groups = human_df[human_df['group'] != 'Overall'].rename(columns={'D_H': 'D_G', 'R_H': 'R_G'})\n",
    "        key_to_ordering = {k: v for k, v in zip(info_df['key'], info_df['option_ordinal'])}\n",
    "        human_groups['ordinal'] = human_groups.apply(lambda x: key_to_ordering[x['qkey']], axis=1)\n",
    "        human_groups = pd.merge(human_groups, human_overall, on='qkey')\n",
    "        human_groups['WD'] = human_groups.apply(lambda x: wasserstein_distance(x['ordinal'], \n",
    "                                                                             x['ordinal'],\n",
    "                                                                             x['D_G'], x['D_O']) / ph.get_max_wd(x['ordinal']), \n",
    "                                              axis=1)\n",
    "        human_groups.to_csv(os.path.join(OUTPUT_DIR, f'{SURVEY_NAME}_baseline.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad01418b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
